package storage

import (
	"context"

	"gitcode.com/eteam/sae.git/types"
	"go.mongodb.org/mongo-driver/v2/bson"
	"go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
)

type TaskStatusMessageDbItem struct {
	MessageId string            `bson:"messageId"`
	TimeStamp int64             `bson:"timeStamp"`
	Role      types.MessageRole `bson:"role"`
	Parts     []types.Part      `bson:"parts"`
}

func (item *TaskStatusMessageDbItem) ToMessage() *types.Message {
	if item == nil {
		return nil
	}
	return &types.Message{
		MessageId: item.MessageId,
		TimeStamp: item.TimeStamp,
		Role:      item.Role,
		Parts:     item.Parts,
		Kind:      types.MessageObjectKind,
	}
}

type TaskStatusDbItem struct {
	State     types.TaskState         `bson:"state"`
	Message   TaskStatusMessageDbItem `bson:"message"`
	TimeStamp int64                   `bson:"timeStamp"`
}

type TaskInfoDbItem struct {
	TaskId  string           `bson:"_id"`
	OwnerId string           `bson:"ownerId"`
	Status  TaskStatusDbItem `bson:"status"`
	SkillId string           `bson:"skillId"`
}

func (item *TaskInfoDbItem) ToTask() *types.Task {
	return &types.Task{
		TaskId:  item.TaskId,
		OwnerId: item.OwnerId,
		Status: types.TaskStatus{
			State:     item.Status.State,
			Message:   item.Status.Message.ToMessage(),
			TimeStamp: item.Status.TimeStamp,
		},
		SkillId:     item.SkillId,
		Messages:    []types.Message{},
		Artifacts:   []types.Artifact{},
		FormDefines: []types.FormDefine{},
	}
}

type _TaskInfoDao struct {
	mc *mongo.Collection
}

func newTaskInfoDao(db *mongo.Database) (*_TaskInfoDao, error) {
	mc := db.Collection("task_info")
	_, err := mc.Indexes().CreateOne(context.Background(), mongo.IndexModel{
		Keys: bson.D{{"status.timeStamp", 1}, {"ownerId", 1}},
	})
	if err != nil {
		return nil, err
	}
	return &_TaskInfoDao{mc: mc}, nil
}

func (dao *_TaskInfoDao) Insert(ctx context.Context, dbItem *TaskInfoDbItem) error {
	_, err := dao.mc.InsertOne(ctx, dbItem)
	return err
}

func (dao *_TaskInfoDao) CountByOwnerId(ctx context.Context, ownerId string) (uint32, error) {
	count, err := dao.mc.CountDocuments(ctx, bson.M{"ownerId": ownerId})
	return uint32(count), err
}

func (dao *_TaskInfoDao) ListByOwnerId(ctx context.Context, ownerId string, offset, limit int64) ([]*TaskInfoDbItem, error) {
	cursor, err := dao.mc.Find(ctx, bson.M{"ownerId": ownerId},
		options.Find().SetSkip(offset).SetLimit(limit).SetSort(bson.M{"status.timeStamp": -1}))
	if err != nil {
		return nil, err
	}
	retList := []*TaskInfoDbItem{}
	err = cursor.All(ctx, &retList)
	return retList, err
}

func (dao *_TaskInfoDao) Get(ctx context.Context, taskId string) (*TaskInfoDbItem, error) {
	retItem := &TaskInfoDbItem{}
	err := dao.mc.FindOne(ctx, bson.M{"_id": taskId}).Decode(retItem)
	return retItem, err
}

func (dao *_TaskInfoDao) UpdateStatus(ctx context.Context, taskId string, status *TaskStatusDbItem) error {
	_, err := dao.mc.UpdateOne(ctx, bson.M{"_id": taskId}, bson.M{"$set": bson.M{"status": status}})
	return err
}
